/*!
  @file feature_extractor.cpp
  @copyright 2013 Kubota Lab. All rights resereved.
*/

#include "mrover_vo/feature_extractor.h"

#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

#include <ros/ros.h>

using std::vector;

namespace vo {

void FeatureExtractor::compute() {
  if (!initCompute()) {
    ROS_ERROR("error in feature tracker");
    return;
  }

  detector_->detect(img_vec_[0], *feature_->kp_vec[0], mask_);
  detector_->detect(img_vec_[1], *feature_->kp_vec[1], mask_);
  descriptor_->compute(img_vec_[0], *feature_->kp_vec[0], *feature_->desc_vec[0]);
  descriptor_->compute(img_vec_[1], *feature_->kp_vec[1], *feature_->desc_vec[1]);

  matcher_->setInputFeature(feature_->kp_vec[0], feature_->desc_vec[0]);
  matcher_->setTargetFeature(feature_->kp_vec[1], feature_->desc_vec[1]);
  matcher_->compute(feature_->matches);

  detectSubpixel();

  triangulation_->setInputFeature(feature_);
  triangulation_->compute(feature_->sp);

  ROS_INFO("kpL %d, kpR %d, match %d", feature_->kp_vec[0]->size(), feature_->kp_vec[1]->size(), feature_->matches->size());
  removeUnpairedFeatures();
}

void FeatureExtractor::detectSubpixel() {
  static cv::Size winSize(5, 5);
  static cv::Size zeroZone(-1, -1);
  static cv::TermCriteria criteria(CV_TERMCRIT_EPS + CV_TERMCRIT_ITER, 40, 0.001 );
  for (size_t d = 0; d < 2; ++d) {
    vector<cv::Point2f> subpix;
    cv::KeyPoint::convert(*feature_->kp_vec[d], subpix);
    cv::cornerSubPix(img_vec_[d], subpix, winSize, zeroZone, criteria);
    cv::KeyPoint::convert(subpix, *feature_->kp_vec[d]);
  }
}

void FeatureExtractor::removeUnpairedFeatures() {
  vector<KeyPointsPtr> tmp_vkp(2);
  tmp_vkp.swap(feature_->kp_vec);
  vector<DescriptorsPtr> tmp_vdesc(2);
  tmp_vdesc.swap(feature_->desc_vec);
 
  // extract matched indices
  vector< vector<int> > indices(2);
  indices[0].resize(tmp_vkp[0]->size(), 0);
  indices[1].resize(tmp_vkp[1]->size(), 0);
  //vector<int> indices1(tmp_vkp[1].size(), 0);
  for (Feature::Matches::iterator match_it = feature_->matches->begin(); 
       match_it != feature_->matches->end(); ++match_it) {
    indices[0][match_it->queryIdx] = 1;
    indices[1][match_it->trainIdx] = 1;
  }

  // compute mapping to new indices
  vector< vector<int> > map(2);
  map[0].resize(tmp_vkp[0]->size(), 0);
  map[1].resize(tmp_vkp[1]->size(), 0);
  for (size_t d = 0; d < 2; ++d) {
    for (size_t i = 0, c = 0; i < indices[d].size(); ++i) {
      if (indices[d][i]) map[d][i] = c++;
    }
  }
  for (Feature::Matches::iterator match_it = feature_->matches->begin(); 
       match_it != feature_->matches->end(); match_it++) {
    match_it->queryIdx = map[0][match_it->queryIdx];
    match_it->trainIdx = map[1][match_it->trainIdx];
  }

  // copy
  feature_->kp_vec[0].reset(new Feature::KeyPoints);
  feature_->kp_vec[1].reset(new Feature::KeyPoints);
  feature_->desc_vec[0].reset(new Feature::Descriptors);
  feature_->desc_vec[1].reset(new Feature::Descriptors);
  feature_->desc_vec[0]->create(cv::Size(tmp_vdesc[0]->cols, 0), tmp_vdesc[0]->type());
  feature_->desc_vec[1]->create(cv::Size(tmp_vdesc[1]->cols, 0), tmp_vdesc[1]->type());
  for (int d = 0; d < 2; ++d) {
    for (size_t i = 0; i < indices[d].size(); ++i) {
      if (indices[d][i]) {
        feature_->kp_vec[d]->push_back(tmp_vkp[d]->at(i));
        feature_->desc_vec[d]->push_back(tmp_vdesc[d]->row(i));
      }
    }
  }
}

}  // namespace vo
